#!/usr/bin/env python3
# 导入必要的库
import argparse
import os
import sys
import time
import grpc
import traceback
import threading

# 添加P4Runtime工具库路径
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 './utils/')
)
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

import struct

H1_IP="10.0.1.1"
H2_IP="10.0.1.2"
H3_IP="10.0.1.3"
H4_IP="10.0.1.4"
H1_MAC="00:00:00:00:01:01"
H2_MAC="00:00:00:00:01:02"
H3_MAC="00:00:00:00:01:03"
H4_MAC="00:00:00:00:01:04"


def writeAcl_tableRules1(p4info_helper, sw, dst_ip):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_table",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, 0xffffffff)
        },
        action_name="MyIngress.drop",
        action_params={},
        priority=10,
        )

    sw.WriteTableEntry(table_entry)

def writeAcl_tableRules2(p4info_helper, sw, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_table",
        match_fields={
            "hdr.udp.dstPort": (port, 65535)
        },
        action_name="MyIngress.drop",
        action_params={},
        priority=20,
        )

    sw.WriteTableEntry(table_entry)

def writeIpv4_lpmRules(p4info_helper, sw, dst_ip, dst_mac, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, 32)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port"   : port,
        })
    sw.WriteTableEntry(table_entry)


def printGrpcError(e):
    """打印gRPC通信错误的详细信息"""
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print(f"({status_code.name})", end=' ')
    traceback_obj = sys.exc_info()[2]
    print(f"[{traceback_obj.tb_frame.f_code.co_filename}:{traceback_obj.tb_lineno}]")




def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建与交换机s1的连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')


        # 发送主控制器仲裁消息
        s1.MasterArbitrationUpdate()

        # 在交换机上安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s1")


        writeIpv4_lpmRules(p4info_helper, s1, H1_IP, H1_MAC, 1)
        writeIpv4_lpmRules(p4info_helper, s1, H2_IP, H2_MAC, 2)
        writeIpv4_lpmRules(p4info_helper, s1, H3_IP, H3_MAC, 3)
        writeIpv4_lpmRules(p4info_helper, s1, H4_IP, H4_MAC, 4)

        writeAcl_tableRules1(p4info_helper, s1, "10.0.1.4")
        writeAcl_tableRules2(p4info_helper, s1, 80)





    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"Unexpected error: {str(e)}")
        traceback.print_exc()

    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/acl.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/acl.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
